package es.upm.dit.gsi.opactv.util;

import java.io.UnsupportedEncodingException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.Set;
import java.util.TimeZone;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.namespace.QName;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPMessage;
import javax.xml.ws.handler.Handler;
import javax.xml.ws.handler.HandlerResolver;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.PortInfo;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;

import org.w3c.dom.Element;
import org.w3c.dom.Node;

public class AwsHandlerResolver implements HandlerResolver {

    private String awsSecretKey;

    public AwsHandlerResolver(String awsSecretKey) {
        this.awsSecretKey = awsSecretKey;
    }

    @SuppressWarnings("unchecked")
    public List<Handler> getHandlerChain(PortInfo portInfo) {
        List<Handler> handlerChain = new ArrayList<Handler>();
        QName serviceQName = portInfo.getServiceName();
        if(serviceQName.getLocalPart().equals("AWSECommerceService")) {
            handlerChain.add(new AwsHandler(awsSecretKey));
        }
        return handlerChain;
    }

    private static class AwsHandler implements SOAPHandler<SOAPMessageContext> {

        private byte[] secretBytes;

        public AwsHandler(String awsSecretKey) {
            secretBytes = stringToUtf8(awsSecretKey);
        }

        public void close(MessageContext messagecontext) {}

	public Set<QName> getHeaders() {
	    return null;
	}

	public boolean handleFault(SOAPMessageContext messagecontext) {
	    return true;
	}

	public boolean handleMessage(SOAPMessageContext messagecontext) {
            Boolean outbound = (Boolean) messagecontext.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);
            if (outbound) {
                try {
                    SOAPMessage soapMessage = messagecontext.getMessage();
                    SOAPBody soapBody = soapMessage.getSOAPBody();
                    Node firstChild = soapBody.getFirstChild();
                    String timeStamp = getTimestamp();
                    String signature = getSignature(firstChild.getLocalName(), timeStamp, secretBytes);
                    appendTextElement(firstChild, "Signature", signature);
                    appendTextElement(firstChild, "Timestamp", timeStamp);
		} catch(SOAPException se) {
                    throw new RuntimeException("SOAPException was thrown.", se);
		}
            }
	    return true;
	}

	private String getSignature(String operation, String timeStamp, byte[] secretBytes) {
            try {
                String toSign = operation + timeStamp;
		byte[] toSignBytes = stringToUtf8(toSign);
		Mac signer = Mac.getInstance("HmacSHA256");
		SecretKeySpec keySpec = new SecretKeySpec(secretBytes, "HmacSHA256");
		signer.init(keySpec);
		signer.update(toSignBytes);
		byte[] signBytes = signer.doFinal();
		String signature = new String(Base64.encode(signBytes));
                return signature;
            } catch(NoSuchAlgorithmException nsae) {
		throw new RuntimeException("NoSuchAlgorithmException was thrown.", nsae);
            } catch(InvalidKeyException ike) {
            	throw new RuntimeException("InvalidKeyException was thrown.", ike);
            }
        }

	private String getTimestamp() {
            Calendar calendar = Calendar.getInstance();
            SimpleDateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss'Z'");
            dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
            return dateFormat.format(calendar.getTime());
        }

        private byte[] stringToUtf8(String source) {
            try {
                return source.getBytes("UTF-8");
            } catch (UnsupportedEncodingException e) {
                // This will never happen.  UTF-8 is always available.
		throw new RuntimeException("getBytes threw an UnsupportedEncodingException", e);
            }
	}

	private void appendTextElement(Node node, String elementName, String elementText) {
            Element element = node.getOwnerDocument().createElement(elementName);
	    element.setTextContent(elementText);
	    node.appendChild(element);
	}

        private static class Base64 {

            private static char[] map1 = new char[64];
            static {
                int i=0;
                for (char c='A'; c<='Z'; c++) map1[i++] = c;
                for (char c='a'; c<='z'; c++) map1[i++] = c;
                for (char c='0'; c<='9'; c++) map1[i++] = c;
                map1[i++] = '+'; map1[i++] = '/';
            }

            private static byte[] map2 = new byte[128];
            static {
                for (int i=0; i<map2.length; i++) map2[i] = -1;
                for (int i=0; i<64; i++) map2[map1[i]] = (byte)i;
            }

            private Base64() {}

            public static char[] encode (byte[] in) {
                return encode(in,in.length);
            }

            public static char[] encode (byte[] in, int iLen) {
                int oDataLen = (iLen*4+2)/3;
                int oLen = ((iLen+2)/3)*4;
                char[] out = new char[oLen];
                int ip = 0;
                int op = 0;
                while (ip < iLen) {
                    int i0 = in[ip++] & 0xff;
                    int i1 = ip < iLen ? in[ip++] & 0xff : 0;
                    int i2 = ip < iLen ? in[ip++] & 0xff : 0;
                    int o0 = i0 >>> 2;
                    int o1 = ((i0 &   3) << 4) | (i1 >>> 4);
                    int o2 = ((i1 & 0xf) << 2) | (i2 >>> 6);
                    int o3 = i2 & 0x3F;
                    out[op++] = map1[o0];
                    out[op++] = map1[o1];
                    out[op] = op < oDataLen ? map1[o2] : '='; op++;
                    out[op] = op < oDataLen ? map1[o3] : '='; op++;
                }
                return out;
            }

            public static byte[] decode (String s) {
                return decode(s.toCharArray());
            }

            public static byte[] decode (char[] in) {
                int iLen = in.length;
                if (iLen%4 != 0) throw new IllegalArgumentException ("Length of Base64 encoded input string is not a multiple of 4.");
                while (iLen > 0 && in[iLen-1] == '=') iLen--;
                int oLen = (iLen*3) / 4;
                byte[] out = new byte[oLen];
                int ip = 0;
                int op = 0;
                while (ip < iLen) {
                    int i0 = in[ip++];
                    int i1 = in[ip++];
                    int i2 = ip < iLen ? in[ip++] : 'A';
                    int i3 = ip < iLen ? in[ip++] : 'A';
                    if (i0 > 127 || i1 > 127 || i2 > 127 || i3 > 127)
                        throw new IllegalArgumentException ("Illegal character in Base64 encoded data.");
                    int b0 = map2[i0];
                    int b1 = map2[i1];
                    int b2 = map2[i2];
                    int b3 = map2[i3];
                    if (b0 < 0 || b1 < 0 || b2 < 0 || b3 < 0)
                        throw new IllegalArgumentException ("Illegal character in Base64 encoded data.");
                    int o0 = ( b0       <<2) | (b1>>>4);
                    int o1 = ((b1 & 0xf)<<4) | (b2>>>2);
                    int o2 = ((b2 &   3)<<6) |  b3;
                    out[op++] = (byte)o0;
                    if (op<oLen) out[op++] = (byte)o1;
                    if (op<oLen) out[op++] = (byte)o2;
                }
                return out;
            }

        }

    }

}
